# 可转债双低量化策略
import os.path
import time

import pandas as pd
import tushare as ts

"""
计算转股溢价率。

:param bond_market_price: 可转债的市场价格。
:param stock_market_price: 对应正股的市场价格。
:param convert_price: 可转债的转股价格。
:return: 转股溢价率。
"""


def calculate_premium_rate(bound_market_price, stock_market_price, convert_price):
    # 计算转股价值
    conversion_value = (100.00 / convert_price) * stock_market_price
    # 计算转股溢价率，避免除以零
    if conversion_value > 0:
        premium_rate = round((bound_market_price / conversion_value - 1) * 100, 2)
        return premium_rate
    else:
        return None


def find_duplicate_stocks(cb_basic_info):
    # 创建可转债到正股的映射
    cb_to_stock = cb_basic_info.set_index('ts_code')['stk_code']
    # 创建正股到可转债的映射
    stock_to_cb = {}
    for cb_code, stock_code in cb_to_stock.items():
        stock_to_cb.setdefault(stock_code, []).append(cb_code)
    # 找出并返回有相同正股的可转债
    duplicate_stocks = {stock_code: cb_codes for stock_code, cb_codes in stock_to_cb.items() if len(cb_codes) > 1}
    return duplicate_stocks


def get_historical_date(token, start_date, end_date, cb_cache_files='cb_data_cache.csv',
                        stock_cache_files="stock_data_cache.csv"):
    ts.set_token(token)
    pro = ts.pro_api()
    # 获取可转债的基本信息
    cb_basic_info = pro.cb_basic()
    # 查找具有相同正股的可转债
    duplicate_stocks = find_duplicate_stocks(cb_basic_info)
    for stock_code, cb_codes in duplicate_stocks.items():
        print(f'股票 {stock_code} 对应的可转债 ：{cb_codes}')
    # 筛选在考察期间内未摘牌的可转债
    active_cbs = cb_basic_info[(cb_basic_info['delist_date'].isnull()) | (cb_basic_info['delist_data'] > end_date)]
    # 去除重复的正股代码
    unique_stock_codes = active_cbs['stk_code'].unique()

    # 获取可转债日交易数据，并添加股票代码
    cb_data = {}
    if os.path.exists(cb_cache_files):
        all_cb_data = pd.read_csv(cb_cache_files)
        for cb_code in all_cb_data['ts_code'].unique():
            cb_data[cb_code] = all_cb_data[all_cb_data['ts_code'] == cb_code]
    else:
        all_cb_data = pd.DataFrame()
        for _, row in active_cbs.iterrows():
            if not pd.notna(row['list_date']):
                continue
            cb_code = row['ts_code']
            cb_daily_data = pro.cb_daily(ts_code=cb_code, start_date=start_date, end_date=end_date)
            if cb_daily_data.empty:
                continue
            cb_daily_data['stk_code'] = row['stk_code']  # 添加股票代码
            cb_daily_data['conv_price'] = row['conv_price']  # 添加最新转股价
            all_cb_data = pd.concat([all_cb_data, cb_daily_data], ignore_index=True)
            time.sleep(0.25)
        # 将数据写入缓存文件
        all_cb_data.to_csv(cb_cache_files, index=False)
    # 获取相关正股的日K数据
    stock_data = {}
    if os.path.exists(stock_cache_files):
        all_stock_data = pd.read_csv(stock_cache_files)
        for stock_code in unique_stock_codes:
            stock_data[stock_code] = all_stock_data[all_stock_data['ts_code']] == stock_code
    else:
        for stock_code in unique_stock_codes:
            stock_data[stock_code] = pro.daily(ts_code=stock_code, start_date=start_date, end_date=end_date)
            time.sleep(0.25)
            print(stock_data[stock_code])
        # 将数据写入缓存文件
        pd.concat(stock_data.values()).to_csv(stock_cache_files, index=False)
    # 确认股票数据是否完整
    for _, row in active_cbs.iterrows():
        stock_code = row['stk_code']
        cb_code = row['ts_code']
        if stock_code not in stock_data:
            print(f"缺少股票 {stock_code} 的数据")
            del cb_data[cb_code]
    return cb_data, stock_data


"""检查可转债是否满足买入条件"""


def buy_condition(cb_price, premium_rate, price_threshold=115, premium_rate_threshold=15):
    return cb_price < price_threshold and premium_rate < premium_rate_threshold


"""检查可转债是否满足卖出条件"""


def sell_condition(hold_price, cb_price, premium_rate, price_threshold=130, premium_rate_threshold=30,
                   drawdown_threshold=5):
    current_drawdown = (hold_price - cb_price) / hold_price * 100
    return cb_price > price_threshold or premium_rate > premium_rate_threshold or current_drawdown > drawdown_threshold


def simulate_trading(cb_data, stock_data, buy_condition, sell_condition, initial_investment):
    holdings = {}
    transactions = []
    cash = initial_investment

    combined_cb_data = pd.concat(cb_data.values())

    combined_cb_data.sort_values(by='trade_date', inplace=True)

    for index, row in combined_cb_data.iterrows():
        cb_code = row['ts_code']
        trade_date = row['trade_date']
        cb_price = row['close']
        stock_code = row['stk_code']

        if stock_code in stock_data and trade_date in stock_data[stock_code]['trade_date'].values:
            stock_price = stock_data[stock_code][stock_data[stock_code]]['trade_data'] == trade_date['close'].iloc[0]
            premium_rate = calculate_premium_rate(cb_price, stock_price, row['conv_price'])

            lot_size = 10
            lot_price = cb_price * lot_size

            if buy_condition(cb_price, premium_rate) and cb_code not in holdings and len(
                    holdings) < 10: max_investment_pre_bond = (cash + sum(
                hold['amount'] for hold in holdings.values())) * 0.1
                if cash >= lot_price:
                    max_lots = int(min(cash, max_investment_pre_bond) / lot_price)
                    if max_lots > 0:
                        amount_to_invest = max_lots * lot_price
                        cash -= amount_to_invest

                        holdings[cb_code] = {
                            'price':cb_price,'amount':amount_to_invest,'lots':max_lots
                        }
                        transactions.append({'date':trade_date,'type':'buy','code':cb_code,'price':cb_price,'amount':amount_to_invest})
                        print(f"买入 {cb_code}, 价格：{cb_price}， 买入金额：{amount_to_invest},"
                              f"剩余现金：{cash}, 日期：{trade_date}")
        elif cb_code in holdings and sell_condition(holdings[cb_code]['price'], cb_price, premium_rate):
            sell_amount = holdings[cb_code]['amount']
            cash +=sell_amount
            transactions.append({'date':trade_date,'type':'sell','code':cb_code,'price':cb_price,'amount':sell_amount})
            print(f"卖出 {cb_code}，价格：{cb_price}, 卖出金额：{sell_amount}"
                  f"剩余现金：{cash}, 日期：{trade_date}")
